import pyaudio
import webrtcvad
import numpy as np



# 参数配置
AUDIO_RATE = 16000       # 采样率：16000（支持8000, 16000, 32000或48000）
CHUNK_SIZE = 480         # 每块大小（30ms，保证为10/20/30ms的倍数）
VAD_MODE = 1             # VAD模式（0-3，数值越小越保守）

# 初始化VAD
vad = webrtcvad.Vad(VAD_MODE)


# 初始化模型
"""
try

pip install modelscope[framework]

or

pip install modelscope[dataset]
"""
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
inference_pipeline = pipeline(
    task=Tasks.auto_speech_recognition,
    model=r'D:\Downloads\SenseVoiceSmall',
    model_revision="master",
    device="cuda:0",
    use_itn=True,
    disable_update=True,
)


# 初始化模型
# from modelscope.pipelines import pipeline
# from modelscope.utils.constant import Tasks
# inference_pipeline = pipeline(
#     task=Tasks.auto_speech_recognition,
#     model=r"D:\Downloads\speech_UniASR_asr_2pass-cn-dialect-16k-vocab8358-tensorflow1-offline")


# 加载模型
# from funasr import AutoModel
# from funasr.utils.postprocess_utils import rich_transcription_postprocess
# model_dir = r'D:\Downloads\SenseVoiceSmall'
# model = AutoModel(
#     model=model_dir,
#     trust_remote_code=True,
#     remote_code="./model.py",  
#     vad_model="fsmn-vad",
#     vad_kwargs={"max_single_segment_time": 30000},
#     device="cuda:0",
#     disable_update=True,
# )

class SpeechDetector:
    def __init__(self, calibration_seconds=2, chunk_duration_ms=30):
        """
        calibration_seconds: 校准背景噪音所需时间（秒）
        chunk_duration_ms: 每块时长（毫秒）
        """
        self.calibration_seconds = calibration_seconds
        self.chunk_duration_ms = chunk_duration_ms
        self.calibrated = False
        self.amplitude_threshold = None

        # 语音缓冲区
        self.speech_buffer = bytearray()

        # 连续帧判决参数（降低短时噪音误判）
        self.speech_state = False         # 当前状态：True为语音，False为无语音
        self.consecutive_speech = 0       # 连续语音帧计数
        self.consecutive_silence = 0      # 连续静音帧计数
        self.required_speech_frames = 2   # 连续3帧语音后确认进入语音状态（约90ms）1帧为30ms
        self.required_silence_frames = 15  # 750ms 连续25帧静音后退出语音状态

    def calibrate(self, stream):
        """
        校准背景噪音：录制固定时长音频，计算平均幅值与标准差，从而设置自适应阈值
        """
        print("开始校准背景噪音，请保持安静...")
        amplitudes = []
        num_frames = int(self.calibration_seconds * (1000 / self.chunk_duration_ms))
        for _ in range(num_frames):
            audio_chunk = stream.read(CHUNK_SIZE, exception_on_overflow=False)
            audio_data = np.frombuffer(audio_chunk, dtype=np.int16)
            amplitudes.append(np.abs(audio_data).mean())
        mean_noise = np.mean(amplitudes)
        std_noise = np.std(amplitudes)
        # 阈值设置为均值加2倍标准差
        self.amplitude_threshold = mean_noise + 2 * std_noise
        print(f"校准完成：噪音均值={mean_noise:.2f}，标准差={std_noise:.2f}，设置阈值={self.amplitude_threshold:.2f}")
        self.calibrated = True

    def analyze_spectrum(self, audio_chunk):
        """
        通过频谱分析检测语音特性：
        1. 对音频块应用汉宁窗后计算 FFT
        2. 统计局部峰值数量（峰值必须超过均值的1.5倍）
        3. 当峰值数量大于等于设定阈值时，认为该块具有语音特征
        """
        audio_data = np.frombuffer(audio_chunk, dtype=np.int16)
        if len(audio_data) == 0:
            return False

        # 使用汉宁窗减少FFT泄露
        window = np.hanning(len(audio_data))
        windowed_data = audio_data * window

        # FFT计算得到频谱（只需正频率部分）
        spectrum = np.abs(np.fft.rfft(windowed_data))
        spectral_mean = np.mean(spectrum)

        # 统计超过均值1.5倍的局部峰值数量
        peak_count = 0
        for i in range(1, len(spectrum) - 1):
            if spectrum[i] > spectrum[i - 1] and spectrum[i] > spectrum[i + 1] and spectrum[i] > spectral_mean * 1.5:
                peak_count += 1

        return peak_count >= 3

    def is_speech(self, audio_chunk):
        """
        综合判断：先通过能量预处理（阈值）过滤低能量数据，再利用VAD和频谱分析判断语音。
        两者结合能有效降低噪音导致的误判。
        """
        amplitude_threshold = self.amplitude_threshold if self.amplitude_threshold is not None else 500
        audio_data = np.frombuffer(audio_chunk, dtype=np.int16)
        amplitude = np.abs(audio_data).mean()
        if amplitude < amplitude_threshold:
            return False

        # VAD检测与频谱检测
        vad_result = vad.is_speech(audio_chunk, AUDIO_RATE)
        spectral_result = self.analyze_spectrum(audio_chunk)

        return vad_result and spectral_result

    def process_chunk(self, audio_chunk):
        """
        对每一块数据进行处理：综合能量检测、VAD、频谱分析，并采用连续帧策略实现状态平滑。
        修改后的逻辑如下：
          - 当检测到语音时（is_speech_chunk 为 True）：
                如果当前未进入语音状态且连续语音帧达到设定阈值，则将状态置为True，并打印提示；
                如果处于语音状态，则不断将当前的 audio_chunk 追加到缓冲区。
          - 当检测为非语音时：
                累计静音帧，若连续静音帧达到设定阈值，则认为语音结束，调用推理模型，
                并清空缓冲区后恢复初始状态。
        """
        is_speech_chunk = self.is_speech(audio_chunk)

        if is_speech_chunk:
            self.consecutive_speech += 1
            self.consecutive_silence = 0
            # 当未进入语音状态且连续语音帧达到阈值时，转为语音状态
            if not self.speech_state and self.consecutive_speech >= self.required_speech_frames:
                self.speech_state = True
                print("Detected Speech")
            # 如果处于语音状态，则持续追加当前块
            if self.speech_state:
                self.speech_buffer.extend(audio_chunk)
        else:
            self.consecutive_silence += 1
            self.consecutive_speech = 0
            # 当处于语音状态并且连续静音帧达到阈值时，结束语音录入
            if self.speech_state and self.consecutive_silence >= self.required_silence_frames:
                
                self.speech_buffer.extend(audio_chunk) # 包含尾部的静音
                
                self.speech_state = False # 结束语音录入

                print("正在识别...", len(self.speech_buffer))
                # 将缓冲区的数据转换为 NumPy 数组
                audio_data = bytes(self.speech_buffer)
                rec_result = self.sound2text(audio_data)
                print("识别结果:", rec_result)

                # 清空缓冲区准备下一段语音
                self.speech_buffer = bytearray()
                print("No speech")
    
    def sound2text(self, audio_data):
        """
        直接输入音频数据进行推理
        """
        # res = model.generate(
        #     input=audio_data,
        #     cache={},
        #     language="zn",  # "zn", "en", "yue", "ja", "ko", "nospeech"
        #     use_itn=True,
        #     batch_size_s=60,
        #     merge_vad=True,  #
        #     merge_length_s=15,
        # )
        # text = rich_transcription_postprocess(res[0]["text"])

        text = inference_pipeline(audio_data)
        return text

def main():
    p = pyaudio.PyAudio()
    stream = p.open(
        format=pyaudio.paInt16,
        channels=1,
        rate=AUDIO_RATE,
        input=True,
        frames_per_buffer=CHUNK_SIZE
    )

    detector = SpeechDetector()
    # 校准背景噪音（建议在程序启动时进行）
    detector.calibrate(stream)

    print("开始监听，请开始说话...(按Ctrl+C停止)")
    try:
        while True:
            audio_chunk = stream.read(CHUNK_SIZE, exception_on_overflow=False)
            detector.process_chunk(audio_chunk)
    except KeyboardInterrupt:
        print("停止监听")
    finally:
        if detector.speech_buffer:
            detector.sound2text(bytes(detector.speech_buffer))
        stream.stop_stream()
        stream.close()
        p.terminate()

if __name__ == "__main__":
    main()
